# -*- coding: utf-8 -*-
# @Author  : longbhu
# @Time    : 2025/2/25 10:29
# @Function:
import glob
import os
import rasterio
import numpy as np
from typing import List, Union
import geopandas as gpd
from rasterio.mask import mask
from osgeo import gdal
from rasterio import open as rio_open
from shapely import intersects
from shapely.geometry import box


# todo 实现将TIF文件中的每个像素值乘以一个给定的参数
def multiply_tif_data_by_factor(
        tif_path: str,
        output_path: str,
        factor: Union[int, float]
) -> np.ndarray[np.float32]:
    """
    将TIF文件中的每个像素值乘以给定的参数，并保存结果到新的TIF文件。

    :param tif_path: 输入的TIF文件路径
    :param output_path: 输出的TIF文件路径
    :param factor: 乘数因子，可以是整数或浮点数
    """
    # 打开TIF文件
    with rasterio.open(tif_path) as src:
        # 读取TIF文件的数据
        data = src.read()
        # 获取TIF文件的元数据
        meta = src.meta.copy()

        # 将数据乘以给定的因子
        data = data * factor

        # 更新元数据中的数据类型（如果因子是浮点数，结果可能是浮点数）
        if isinstance(factor, float):
            meta.update(dtype=np.float32)
        else:
            meta.update(dtype=data.dtype)

    # 保存结果到新的TIF文件
    with rasterio.open(output_path, 'w', **meta) as dst:
        dst.write(data)

    print(f"处理完成，结果已保存到 {output_path}")
    return data


def dev_multiply_tif_data_by_factor(input_path, output_path, factor, nodata=-9999):
    """
    将TIFF文件数据乘以指定系数，并处理无效值

    参数:
    input_path: 输入tif路径
    output_path: 输出tif路径
    factor: 乘数
    nodata: 无效值标记（自动识别原始文件的nodata值）
    """
    with rio_open(input_path) as src:
        # 读取数据和元数据
        data = src.read(1)
        meta = src.meta.copy()

        # 获取原始无效值
        original_nodata = src.nodata

        # 处理无效值
        data = np.where(data == original_nodata, np.nan, data)

        # 执行乘法运算（自动忽略nan值）
        result = data * factor

        # 恢复无效值标记
        result = np.where(np.isnan(result), original_nodata, result)

        # 更新元数据
        meta.update({
            'driver': 'GTiff',
            'count': 1,
            'dtype': result.dtype,
            'nodata': original_nodata
        })

        # 写入输出文件
        with rio_open(output_path, 'w', **meta) as dst:
            dst.write(result, 1)


def sum_tif_files(
        input_path: Union[str, List[str]],
        output_filename: str = "sum_result.tif"
) -> np.ndarray[np.float32]:
    """
    将多个TIF文件中的像素值加和，并保存结果到新的TIF文件。
    如果输入的是文件夹路径，则合并文件夹中的所有TIF文件。

    :param input_path: 输入的TIF文件路径列表或文件夹路径
    :param output_filename: 输出的TIF文件名，默认为 "sum_result.tif"
    """
    # 如果输入是文件夹路径，获取文件夹中的所有TIF文件
    if isinstance(input_path, str) and os.path.isdir(input_path):
        tif_paths = [os.path.join(input_path, f) for f in os.listdir(input_path) if f.endswith('.tif')]
        output_path = os.path.join(input_path, output_filename)
    elif isinstance(input_path, list):
        tif_paths = input_path
        output_path = output_filename
    else:
        raise ValueError("输入必须是文件夹路径或TIF文件路径列表！")

    # 检查是否至少有一个TIF文件
    if not tif_paths:
        raise ValueError("未找到任何TIF文件！")

    # 初始化结果数据和元数据
    result_data = None
    meta = None

    # 遍历所有TIF文件
    for tif_path in tif_paths:
        with rasterio.open(tif_path) as src:
            # 如果是第一个文件，初始化结果数据和元数据
            if result_data is None:
                result_data = src.read().astype(np.float32)  # 使用float32避免溢出
                meta = src.meta.copy()
            else:
                # 检查当前文件的大小和数据类型是否与第一个文件一致
                if src.shape != result_data.shape[1:]:
                    raise ValueError(f"文件 {tif_path} 的大小与第一个文件不一致！")
                if src.dtypes[0] != meta['dtype']:
                    raise ValueError(f"文件 {tif_path} 的数据类型与第一个文件不一致！")

                # 将当前文件的数据加到结果数据中
                result_data += src.read().astype(np.float32)

    # 更新元数据中的数据类型
    meta.update(dtype=np.float32)

    # 保存结果到新的TIF文件
    with rasterio.open(output_path, 'w', **meta) as dst:
        dst.write(result_data)

    print(f"加和完成，结果已保存到 {output_path}")
    return result_data


# 示例 1：输入文件夹路径
# 合并文件夹中的所有TIF文件，并将结果保存到同一文件夹中
# sum_tif_files(r'E:\AVIM_output\30m\jiangxi\nep', output_filename='nep_sum_result.tif')

# 示例 2：输入TIF文件路径列表
# 合并指定的TIF文件，并将结果保存到指定路径
# tif_files = ['/path/to/file1.tif', '/path/to/file2.tif', '/path/to/file3.tif']
# sum_tif_files(tif_files, output_filename='/path/to/output_folder/sum_result.tif')


def sum_and_average_tif_pixels(tif_path: str) -> tuple:
    """
    计算TIF文件中所有像素值的总和和平均值。

    :param tif_path: 输入的TIF文件路径
    :return: 一个元组，包含像素值的总和和平均值
    """
    # 打开TIF文件
    with rasterio.open(tif_path) as src:
        # 读取TIF文件的数据
        data = src.read()

        # 检查是否存在无效值（nodata）
        if src.nodata is not None:
            # 将无效值替换为0，避免影响加和和平均计算
            data[data == src.nodata] = 0

        # 计算像素值的总和
        total_sum = np.sum(data)

        # 计算像素值的平均值
        total_pixels = np.count_nonzero(data != src.nodata) if src.nodata is not None else data.size  # total_pixels单位：元
        average_value = total_sum / total_pixels / 900  # 每个点是900m^2  average_value单位：元/m^2

    return total_sum, average_value


# 示例：计算TIF文件中所有像素值的总和和平均值
# tif_path = r'F:\code\dev\calc-gep-regulate-cqc\calc_v1\output\Qwp\sum_result.tif'
# total_sum, average_value = sum_and_average_tif_pixels(tif_path)
#
# print(f"像素值总和: {total_sum}")
# print(f"像素值平均值: {average_value}")

# 根据shp文件裁剪tif文件

def tif_mask_by_shapefile(
        tif_path: str,
        shapefile_folder: str,
        output_folder: str = './masked_output/',
        city_name: str = 'name'
) -> None:
    """
    根据市级 Shapefile 裁剪 TIFF 文件，并将结果保存到指定输出目录。

    参数:
        tif_path (str): 输入的 TIFF 文件路径。
        shapefile_folder (str): 包含 Shapefile 的文件夹路径。
        output_folder (str): 输出目录路径。
        city_name (str): shp文件中的city_name字段的列名。
    """
    # 创建输出目录（如果不存在）
    os.makedirs(output_folder, exist_ok=True)

    # 获取 TIFF 文件名（去掉路径和扩展名） 水源涵养_江西   用split按照"_"切分得到第一个字段
    tif_name_base = os.path.splitext(os.path.basename(tif_path))[0].split('_')[0]
    # 读取 TIFF 文件并获取其 CRS
    with rasterio.open(tif_path) as src:
        tif_bounds = src.bounds
        tif_crs = src.crs  # 获取 TIFF 的 CRS

        # 将栅格边界转换为 shapely 的 Polygon 对象
        tif_polygon = box(*tif_bounds)
    # 读取市级 Shapefile
    shapefiles = glob.glob(os.path.join(shapefile_folder, '*.shp'))

    if not shapefiles:
        print(f"错误: 在 {shapefile_folder} 中未找到任何 Shapefile 文件。")
        return  # 提前返回

    for shp_path in shapefiles:
        # 读取每个 Shapefile
        import chardet
        # 检测编码
        with open(shp_path.replace('.shp', '.dbf'), 'rb') as f:
            result = chardet.detect(f.read())
        detected_encoding = result['encoding']
        print("检测编码为：", detected_encoding)
        city_gdf = gpd.read_file(shp_path, encoding=detected_encoding)

        shp_bounds = city_gdf.total_bounds

        if city_gdf.empty:
            print(f"警告: Shapefile {shp_path} 为空。")
            continue  # 跳过空的 Shapefile

        # 确保坐标系一致
        if city_gdf.crs != tif_crs:
            print(f"警告: Shapefile {shp_path} 的坐标系为 ({city_gdf.crs}) 。")
            print(f"警告: tif {tif_path} 的坐标系为 ({tif_crs}) 。")
            if tif_crs:
                city_gdf = city_gdf.to_crs(tif_crs)
                print("坐标一致")
                print(f"检查: Shapefile {shp_path} 的坐标系为 ({city_gdf.crs}) 。")
                print(f"检查: tif {tif_path} 的坐标系为 ({tif_crs}) 。")
            else:
                src = city_gdf.to_crs(city_gdf.crs)
                tif_crs = src.crs  # 获取 TIFF 的 CRS
                print(f"检查: Shapefile {shp_path} 的坐标系为 ({city_gdf.crs}) 。")
                print(f"检查: tif {tif_path} 的坐标系为 ({tif_crs}) 。")

        # 将矢量边界转换为 shapely 的 Polygon 对象
        shp_polygon = box(*shp_bounds)

        # 空间交集验证
        if not intersects(tif_polygon, shp_polygon):
            raise ValueError("矢量与栅格数据空间范围无重叠，请检查数据源区域匹配性")
        else:
            # 执行mask操作
            # 遍历每个市
            for index, row in city_gdf.iterrows():
                city_geom = [row['geometry']]  # geometry
                shp_city_name = row[city_name]

                # 动态命名输出文件
                output_filename = f"{tif_name_base}_{shp_city_name}.tif"

                # 裁剪 TIFF 文件
                with rasterio.open(tif_path) as src:  # 在这里打开数据集
                    out_image, out_transform = mask(src, city_geom, crop=True)

                    # 更新 metadata
                    out_meta = src.meta.copy()
                    out_meta.update({
                        "height": out_image.shape[1],
                        "width": out_image.shape[2],
                        "transform": out_transform,
                        "crs":"EPSG:4326"
                    })

                    # 保存裁剪后的 TIFF 文件
                    output_path = os.path.join(output_folder, output_filename)

                    with rasterio.open(output_path, 'w', **out_meta) as dest:
                        dest.write(out_image)

                print(f"masked已保存: {output_path}")


# tif_path = r'F:\test_files\插值验证程序的输入数据\fortran插值\et\et_2023.tif'
# shapefile_folder = r'E:\shp\shangluo_shp'
# # output_folder = './output/' # 可选
#
# tif_mask_by_shapefile(tif_path, shapefile_folder,city_name = 'Name') # xiang_name

def process_all_tifs(tif_folder, shapefile_folder, output_folder):
    """
    遍历指定文件夹中的所有 TIFF 文件，并对每个文件调用 tif_mask_by_shapefile 函数。

    参数:
        tif_folder (str): 包含 TIFF 文件的文件夹路径。
        shapefile_folder (str): 包含 Shapefile 的文件夹路径。
        output_folder (str): 输出目录路径。
    """
    # 获取文件夹中所有 TIFF 文件
    tif_files = [f for f in os.listdir(tif_folder) if f.endswith('.tif')]

    # 遍历每个 TIFF 文件，调用裁剪函数
    for tif_file in tif_files:
        tif_path = os.path.join(tif_folder, tif_file)
        print(f"处理文件: {tif_path}")
        tif_mask_by_shapefile(tif_path, shapefile_folder, output_folder)


def merge_tif_files_for_one(input_path, merge_name):
    """
    合并文件夹下的所有tif文件为一个文件。
    输出文件保存在输入的文件夹下创建的merged文件夹中
    参数:
    input_path -- 输入文件夹的路径，包含所有tif文件。
    merge_name -- 输出合并文件的名称。
    """
    # 获取文件夹下的所有tif文件
    input_files = glob.glob(os.path.join(input_path, '*.tif'))
    print(f"查找到 {len(input_files)} 个tif文件。")

    if not input_files:
        print("没有找到任何tif文件。")
        return

    # 存储有效的tif文件
    inputrasfile = []

    for fn in input_files:
        infn = gdal.Open(fn, gdal.GA_ReadOnly)
        if infn is None:
            print(f"文件 {fn} 无法打开，跳过此文件。")
            continue
        input_proj1 = infn.GetProjection()
        if input_proj1 is None:
            print(f"文件 {fn} 没有有效的投影信息，跳过此文件。")
            continue
        inputrasfile.append(infn)

    # 如果没有有效的输入文件，退出函数
    if not inputrasfile:
        print("没有有效的输入文件，无法进行合并。")
        return

    # 确保输出文件夹存在
    output_file = os.path.join(input_path, 'merged', merge_name + '_merged.tif')
    os.makedirs(os.path.dirname(output_file), exist_ok=True)

    # 使用gdal.Warp进行合并
    # 设置输出参数，可以指定投影信息、目标分辨率、合并方法等
    options = gdal.WarpOptions(format='GTiff')
    gdal.Warp(output_file, inputrasfile, options=options)

    print(f"合并完成，文件保存为: {output_file}")

# 示例调用
# merge_tif_files_for_one(r'E:\AVIM_output\30m\shaanxi\nep', 'shaanxi')
